import pygame
import time
import animate_curve


anima = animate_curve.Animation_MP("", ["F", "V", "X"])
pygame.init()
screen_width = 800
screen_height = 600
target_x = 150
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("Ball with Damping")
mass = 0.6  # 质量
ball_radius = 10  # 半径
damping = 4.5 # 阻尼系数
spring_constant = 25  # 弹簧常数
initial_v = 0    # 初始速度
initial_x = 0  # 初始位置
y = 100
time_step = 1 /60


def calculate_position_and_velocity(time_step, mass, c, k, f, x, v, target_x=0):
   dx = x - target_x
   f_ = k * dx - c * v
   v_ = v + (f / mass) * time_step
   x_ = x - v * time_step
#    print(f, f_)
   return f_, v_, x_

def draw_ball(screen, x, y, ball_radius):
   pygame.draw.circle(screen, (255, 255, 255), (int(x), int(y)), ball_radius)


def main():
   global target_x, y
   x = initial_x
   v = initial_v
   f = initial_x * spring_constant - damping * initial_v
   running = True
   while running:
      # 计算小球的位置和速度
      # time_step = 0.1
      f, v, x = calculate_position_and_velocity(time_step, mass, damping, spring_constant, 
                                                         f, x, v, target_x)
      # print(f"f={f:.2f}\tv={v:.2f}\tx={x:.2f}")
      # 更新屏幕上的小球
      anima.update([f, v, x])
      screen.fill((0, 0, 0))
      draw_ball(screen, x, y, ball_radius)
      time.sleep(time_step)
      pygame.display.flip()
      for event in pygame.event.get():
         if event.type == pygame.QUIT:
               running = False
               pygame.quit()
         if event.type == pygame.MOUSEBUTTONDOWN:
            target_x = event.pos[0]
            y = event.pos[1]
         if event.type == pygame.MOUSEMOTION:
            pygame.display.set_caption(f"{event.pos[0]}, {event.pos[1]}")

if __name__ == "__main__":
   anima.start()
   main()
   anima.quit()